"""Utilities for working on json files efficiently,
    including customized encoder and decoder that can support more data types,
    including np.ndarray and datetime.date, datetime.time, datetime.datetime
"""

import datetime
import json
from shapely.geometry.polygon import Polygon
import numpy as np

DATE_FORMAT = "%Y-%m-%d"
TIME_FORMAT = "%H:%M:%S.%f"
DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%f"


class SkywalkerJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime.datetime):
            return {
                "_type": "datetime",
                "value": obj.strftime(DATETIME_FORMAT),
            }
        elif isinstance(obj, datetime.date):
            return {
                "_type": "date",
                "value": obj.strftime(DATE_FORMAT),
            }
        elif isinstance(obj, datetime.time):
            return {
                "_type": "time",
                "value": obj.strftime(TIME_FORMAT),
            }
        elif isinstance(obj, np.ndarray):
            return {"_type": "ndarray", "value": obj.tolist()}
        elif isinstance(obj, Polygon):
            return obj.__geo_interface__
        return super(SkywalkerJSONEncoder, self).default(obj)


class SkywalkerJSONDecoder(json.JSONDecoder):
    def __init__(self, *args, **kwargs):
        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)

    def object_hook(self, obj):
        if "_type" not in obj:
            return obj
        value_type = obj["_type"]
        if value_type == "datetime":
            return datetime.datetime.strptime(obj["value"], DATETIME_FORMAT)
        elif value_type == "date":
            return datetime.datetime.strptime(obj["value"], DATE_FORMAT).date()
        elif value_type == "time":
            return datetime.datetime.strptime(obj["value"], TIME_FORMAT).time()
        elif value_type == "ndarray":
            return np.array(obj["value"])
        return obj


if __name__ == "__main__":
    data = {
        "name": "Silent Bob",
        "dt": datetime.datetime(2013, 11, 11, 10, 40, 32),
        "acquired_date": datetime.date(2019, 1, 1),
        "fake_time": datetime.time(9, 12, 14, 29),
        "raster_data": np.random.randint(0, 255, size=(4, 4)),
        "list_of_dates": [
            datetime.datetime(2013, 11, 11, 10, 40, 32),
            datetime.datetime(2014, 11, 11, 10, 40, 32),
            datetime.datetime(2013, 12, 12, 10, 40, 32),
        ],
        "list_of_arrays": [
            np.arange(0, 1.0, 0.1),
            np.random.randint(0, 16383, size=(5, 5), dtype=np.uint16),
        ],
    }

    str_data = json.dumps(data, cls=SkywalkerJSONEncoder, indent=2)
    print(str_data)

    back_data = json.loads(str_data, cls=SkywalkerJSONDecoder)
    print(back_data)
